#aaabbbccc
import sqlparse
from sqlparse.sql import Identifier, IdentifierList, Where
from sqlparse.tokens import Keyword, DML
import matplotlib.pyplot as plt
import pandas as pd
import sqlite3
from io import StringIO
import re
import time
import subprocess

class SQLOptimizer:
    def __init__(self):
        self.optimization_rules = [
            self.rewrite_select_star,
            self.add_missing_indexes,
            self.simplify_expressions,
            self.rewrite_subqueries,
            self.optimize_joins,
            self.add_limit_to_update_delete
        ]
    
    def optimize(self, sql):
        parsed = sqlparse.parse(sql)[0]
        original_sql = sql
        
        # 应用所有优化规则
        for rule in self.optimization_rules:
            parsed = rule(parsed)
        
        optimized_sql = str(parsed)
        
        return {
            "original": original_sql,
            "optimized": optimized_sql,
            "suggestions": self.generate_suggestions(parsed)
        }
    
    def rewrite_select_star(self, parsed):
        """将SELECT *替换为具体列名 [5,8](@ref)"""
        for token in parsed.tokens:
            if token.ttype is Keyword and token.value.upper() == 'SELECT':
                idx = parsed.tokens.index(token)
                next_token = parsed.tokens[idx + 1]
                
                if isinstance(next_token, IdentifierList) and '*' in next_token.value:
                    # 在实际应用中应查询数据库元数据
                    parsed.tokens[idx + 1].value = 'id, name, date_created'
                elif isinstance(next_token, Identifier) and next_token.value == '*':
                    parsed.tokens[idx + 1].value = 'id, name, date_created'
        return parsed
    
    def add_missing_indexes(self, parsed):
        """识别需要索引的列并生成建议 [2,7](@ref)"""
        where_clause = None
        for token in parsed.tokens:
            if isinstance(token, Where):
                where_clause = token
                break
        
        if where_clause:
            # 在实际应用中应检查索引是否存在
            columns = self.extract_columns(where_clause)
            for col in columns:
                parsed.tokens.append(
                    sqlparse.parse(f"-- SUGGESTION: Add index for {col}")[0]
                )
        return parsed
    
    def simplify_expressions(self, parsed):
        """简化冗余表达式 [8](@ref)"""
        sql = str(parsed)
        # 简化BETWEEN表达式
        sql = re.sub(r"(\w+)\s+>\s+(\d+)\s+AND\s+\1\s+<\s+(\d+)", 
                    r"\1 BETWEEN \2 AND \3", sql, flags=re.IGNORECASE)
        # 简化双重否定
        sql = re.sub(r"NOT\s+NOT\s+(\w+)", r"\1", sql, flags=re.IGNORECASE)
        return sqlparse.parse(sql)[0]
    
    def rewrite_subqueries(self, parsed):
        """将子查询转换为JOIN [5](@ref)"""
        sql = str(parsed)
        # 将IN子查询转为JOIN
        sql = re.sub(
            r"WHERE\s+(\w+)\s+IN\s*\(\s*SELECT\s+(\w+)\s+FROM\s+(\w+)\s*\)",
            r"JOIN \3 ON \1 = \3.\2 WHERE", 
            sql, flags=re.IGNORECASE
        )
        return sqlparse.parse(sql)[0]
    
    def optimize_joins(self, parsed):
        """优化JOIN顺序 [8](@ref)"""
        # 在实际应用中应分析表大小和数据分布
        return parsed
    
    def add_limit_to_update_delete(self, parsed):
        """为UPDATE/DELETE添加LIMIT [8](@ref)"""
        first_token = parsed.token_first()
        if first_token.ttype is DML and first_token.value.upper() in ('UPDATE', 'DELETE'):
            has_limit = any('LIMIT' in t.value.upper() for t in parsed.tokens)
            if not has_limit:
                parsed.tokens.append(
                    sqlparse.parse(" LIMIT 1000 -- Added for safety")[0]
                )
        return parsed
    
    def extract_columns(self, clause):
        """提取WHERE子句中的列名"""
        columns = set()
        for token in clause.tokens:
            if isinstance(token, Identifier):
                columns.add(token.value.split('.')[-1])
        return columns
    
    def generate_suggestions(self, parsed):
        """生成优化建议 [6,8](@ref)"""
        suggestions = []
        sql = str(parsed).upper()
        
        if 'SELECT *' in sql:
            suggestions.append("避免使用SELECT *，明确指定需要的列")
        
        if 'CROSS JOIN' in sql:
            suggestions.append("检查CROSS JOIN是否必要，可能产生笛卡尔积")
        
        if re.search(r"WHERE\s+[\w_]+\s+LIKE\s+'%.*%'", sql):
            suggestions.append("前导通配符LIKE查询无法使用索引，考虑全文搜索")
        
        if 'ORDER BY RAND()' in sql:
            suggestions.append("ORDER BY RAND()性能差，考虑其他随机抽样方法")
        
        return suggestions


class Visualizer:
    @staticmethod
    def compare_execution_plans(original, optimized, db_path=':memory:'):
        """可视化执行计划对比 [6](@ref)"""
        conn = sqlite3.connect(db_path)
        
        # 创建示例数据
        Visualizer.create_sample_data(conn)
        
        # 获取执行计划
        original_plan = Visualizer.get_execution_plan(conn, original)
        optimized_plan = Visualizer.get_execution_plan(conn, optimized)
        
        # 执行时间对比
        original_time = Visualizer.measure_execution_time(conn, original)
        optimized_time = Visualizer.measure_execution_time(conn, optimized)
        
        # 生成对比图表
        fig, ax = plt.subplots(2, 1, figsize=(12, 10))
        
        # 执行计划步骤对比
        steps = list(range(1, max(len(original_plan), len(optimized_plan)) + 1))
        ax[0].plot(steps[:len(original_plan)], [r['order'] for r in original_plan], 
                  'ro-', label='Original')
        ax[0].plot(steps[:len(optimized_plan)], [r['order'] for r in optimized_plan], 
                  'go-', label='Optimized')
        ax[0].set_title('Execution Plan Steps Comparison')
        ax[0].set_xlabel('Step')
        ax[0].set_ylabel('Complexity')
        ax[0].legend()
        ax[0].grid(True)
        
        # 执行时间对比
        time_data = {'Original': original_time, 'Optimized': optimized_time}
        ax[1].bar(time_data.keys(), time_data.values(), color=['red', 'green'])
        ax[1].set_title('Execution Time Comparison (ms)')
        ax[1].set_ylabel('Time (ms)')
        
        for i, v in enumerate(time_data.values()):
            ax[1].text(i, v + 0.5, str(v), ha='center')
        
        plt.tight_layout()
        
        # 保存并显示图表
        plt.savefig('comparison.png')
        plt.show()
        
        # 返回执行计划详情
        return {
            'original_plan': original_plan,
            'optimized_plan': optimized_plan,
            'original_time': original_time,
            'optimized_time': optimized_time
        }
    
    @staticmethod
    def create_sample_data(conn):
        """创建示例数据库"""
        conn.executescript("""
        CREATE TABLE users (
            id INTEGER PRIMARY KEY,
            name TEXT,
            email TEXT,
            created_at TIMESTAMP
        );
        
        CREATE TABLE orders (
            id INTEGER PRIMARY KEY,
            user_id INTEGER,
            amount REAL,
            status TEXT,
            created_at TIMESTAMP,
            FOREIGN KEY(user_id) REFERENCES users(id)
        );
        
        INSERT INTO users (name, email, created_at) VALUES
            ('Alice', 'alice@example.com', '2023-01-01'),
            ('Bob', 'bob@example.com', '2023-01-02'),
            ('Charlie', 'charlie@example.com', '2023-01-03');
        
        INSERT INTO orders (user_id, amount, status, created_at) VALUES
            (1, 100.0, 'completed', '2023-02-01'),
            (1, 50.0, 'pending', '2023-02-02'),
            (2, 200.0, 'completed', '2023-02-03'),
            (3, 75.0, 'cancelled', '2023-02-04');
        """)
    
    @staticmethod
    def get_execution_plan(conn, sql):
        """获取SQL执行计划 [2](@ref)"""
        plan = []
        cursor = conn.execute(f"EXPLAIN QUERY PLAN {sql}")
        for row in cursor:
            plan.append({
                'order': row[0],
                'from': row[1],
                'detail': row[2],
                'sql': row[3]
            })
        return plan
    
    @staticmethod
    def measure_execution_time(conn, sql, runs=10):
        """测量SQL执行时间 [6](@ref)"""
        total_time = 0
        for _ in range(runs):
            start = time.perf_counter()
            conn.execute(sql)
            end = time.perf_counter()
            total_time += (end - start) * 1000  # ms
        return total_time / runs


def main():
    print("SQL Optimizer Tool")
    print("=" * 40)
    
    # 示例SQL或用户输入
    sample_sql = """
    SELECT * FROM users
    WHERE id IN (SELECT user_id FROM orders WHERE status = 'completed')
    ORDER BY RAND();
    """
    
    user_sql = input("Enter your SQL statement (or press Enter to use example):\n> ")
    if not user_sql.strip():
        user_sql = sample_sql
        print(f"Using example SQL:\n{user_sql}")
    
    # 优化SQL
    optimizer = SQLOptimizer()
    result = optimizer.optimize(user_sql)
    
    print("\nOptimization Results")
    print("=" * 40)
    print(f"Original SQL:\n{result['original']}\n")
    print(f"Optimized SQL:\n{result['optimized']}\n")
    
    if result['suggestions']:
        print("Optimization Suggestions:")
        for i, suggestion in enumerate(result['suggestions'], 1):
            print(f"{i}. {suggestion}")
    
    # 执行计划可视化
    print("\nGenerating execution plan comparison...")
    plan_data = Visualizer.compare_execution_plans(
        result['original'], result['optimized']
    )
    
    print("\nPerformance Improvement Summary")
    print("=" * 40)
    improvement = (plan_data['original_time'] - plan_data['optimized_time']) / plan_data['original_time'] * 100
    print(f"Original execution time: {plan_data['original_time']:.2f} ms")
    print(f"Optimized execution time: {plan_data['optimized_time']:.2f} ms")
    print(f"Improvement: {improvement:.1f}%")
    
    # 在浏览器中打开可视化报告
    html_report = generate_html_report(result, plan_data)
    with open('report.html', 'w') as f:
        f.write(html_report)
    
    print("\nOpening detailed report in browser...")
    subprocess.run(['open', 'report.html'])  # macOS
    # subprocess.run(['xdg-open', 'report.html'])  # Linux
    # subprocess.run(['start', 'report.html'], shell=True)  # Windows


def generate_html_report(optimization_result, plan_data):
    """生成HTML格式的详细报告"""
    return f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>SQL Optimization Report</title>
        <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 20px; }}
            .container {{ max-width: 1200px; margin: 0 auto; }}
            .card {{ background: #f9f9f9; border: 1px solid #ddd; padding: 20px; margin-bottom: 20px; }}
            .comparison {{ display: flex; }}
            .sql-panel {{ flex: 1; padding: 10px; }}
            .original {{ border-right: 1px solid #ccc; }}
            .suggestion {{ color: #d35400; margin: 5px 0; }}
            .improvement {{ color: #27ae60; font-weight: bold; }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>SQL Optimization Report</h1>
            
            <div class="card">
                <h2>SQL Comparison</h2>
                <div class="comparison">
                    <div class="sql-panel original">
                        <h3>Original SQL</h3>
                        <pre>{optimization_result['original']}</pre>
                    </div>
                    <div class="sql-panel">
                        <h3>Optimized SQL</h3>
                        <pre>{optimization_result['optimized']}</pre>
                    </div>
                </div>
            </div>
            
            <div class="card">
                <h2>Optimization Suggestions</h2>
                <ul>
                    {''.join(f'<li class="suggestion">{s}</li>' for s in optimization_result['suggestions'])}
                </ul>
            </div>
            
            <div class="card">
                <h2>Performance Comparison</h2>
                <div class="comparison">
                    <div class="sql-panel">
                        <h3>Execution Time</h3>
                        <p>Original: {plan_data['original_time']:.2f} ms</p>
                        <p>Optimized: <span class="improvement">{plan_data['optimized_time']:.2f} ms</span></p>
                        <p>Improvement: <span class="improvement">
                            {(plan_data['original_time'] - plan_data['optimized_time']) / plan_data['original_time'] * 100:.1f}%
                        </span></p>
                    </div>
                    <div class="sql-panel">
                        <h3>Execution Plan Complexity</h3>
                        <canvas id="planChart" width="400" height="200"></canvas>
                    </div>
                </div>
            </div>
            
            <div class="card">
                <h2>Execution Plan Details</h2>
                <div class="comparison">
                    <div class="sql-panel original">
                        <h3>Original Plan</h3>
                        <pre>{format_plan(plan_data['original_plan'])}</pre>
                    </div>
                    <div class="sql-panel">
                        <h3>Optimized Plan</h3>
                        <pre>{format_plan(plan_data['optimized_plan'])}</pre>
                    </div>
                </div>
            </div>
        </div>
        
        <script>
            // 执行计划复杂度图表
            const ctx = document.getElementById('planChart').getContext('2d');
            const chart = new Chart(ctx, {{
                type: 'bar',
                data: {{
                    labels: ['Original', 'Optimized'],
                    datasets: [{{
                        label: 'Execution Steps',
                        data: [{len(plan_data['original_plan'])}, {len(plan_data['optimized_plan'])}],
                        backgroundColor: ['rgba(255, 99, 132, 0.5)', 'rgba(75, 192, 192, 0.5)'],
                        borderColor: ['rgb(255, 99, 132)', 'rgb(75, 192, 192)'],
                        borderWidth: 1
                    }}]
                }},
                options: {{
                    responsive: true,
                    scales: {{
                        y: {{
                            beginAtZero: true,
                            title: {{
                                display: true,
                                text: 'Number of Steps'
                            }}
                        }}
                    }}
                }}
            }});
        </script>
    </body>
    </html>
    """


def format_plan(plan):
    """格式化执行计划输出"""
    buffer = StringIO()
    for step in plan:
        buffer.write(f"Step {step['order']}: {step['detail']}\n")
        buffer.write(f"  From: {step['from']}\n")
        if step['sql']:
            buffer.write(f"  SQL: {step['sql']}\n")
        buffer.write("\n")
    return buffer.getvalue()


if __name__ == "__main__":
    main()